import re
import re
from datetime import datetime
from typing import Union, List

import pandas as pd

from scraper.option_risk_analysis_em import fetch_all_option_risk
from scraper.option_value_analysis_em import fetch_all_option_value

'''
  采集期权实时的value和risk数据
  对value和risk进行分析
  :return:
'''
# 虚值度阈值映射表 - 对应标的代码的虚值度阈值(%)
VIR_VALUE_DEGREE_THRESHOLD_DICT = {
    510050: 5.0,  # 上证50ETF
    510300: 10.0,  # 沪深300ETF
    159919: 10.0,  # 沪深300ETF
    510500: 15.0,  # 中证500ETF
    159922: 15.0,  # 中证500ETF
    588000: 20.0,  # 科创50ETF
    588080: 20.0,  # 科创50ETF
    159901: 20.0,  # 深证100ETF
    159915: 20.0,  # 创业板ETF
}
# 年化收益率阈值
ANNUALIZED_RETURN_THRESHOLD_DICT = {
    "认购": 16.0,
    "认沽": 30.0,
}
# 期权 delta参数的阈值
DELTA_THRESHOLD_DICT = {
    "认购": 0.06,
    "认沽": 0.20,
}


def fetch_option_risk_value() -> pd.DataFrame:  # type: ignore
    '''
    拉取期权的value和risk数据。
    :return:
    '''

    risk_df = fetch_all_option_risk()
    risk_df.to_csv("../result/option_risk_analysis_em.csv", index=False, encoding='utf-8-sig')

    value_df = fetch_all_option_value()
    value_df.to_csv("../result/option_value_analysis_em.csv", index=False, encoding='utf-8-sig')

    combined_df = pd.merge(
        risk_df,
        value_df,
        on=["期权代码", "期权名称", "到期日"],
        how="outer",
        indicator=True
    )

    print("\n合并结果统计:")
    print(f"总行数: {len(combined_df)}")
    print("合并情况:")
    print(combined_df["_merge"].value_counts())

    combined_df.columns = ["期权代码",
                           "期权名称",
                           "最新价",
                           "涨跌幅",
                           "杠杆比率",
                           "实际杠杆比率",
                           "Delta",
                           "Gamma",
                           "Vega",
                           "Rho",
                           "Theta",
                           "到期日",
                           "最新价_y",
                           "时间价值",
                           "内在价值",
                           "隐含波动率",
                           "理论价格",
                           "标的代码",
                           "标的名称",
                           "标的最新价",
                           "标的近一年波动率",
                           "_merge"
                           ]
    combined_df = combined_df[
        [
            "期权代码",
            "期权名称",
            "最新价",
            "时间价值",
            "内在价值",
            "隐含波动率",
            "理论价格",
            "标的代码",
            "标的名称",
            "标的最新价",
            "标的近一年波动率",
            "到期日",
            "涨跌幅",
            "杠杆比率",
            "实际杠杆比率",
            "Delta",
            "Gamma",
            "Vega",
            "Rho",
            "Theta",
        ]
    ]
    combined_df["期权代码"] = pd.to_numeric(combined_df["期权代码"], errors="coerce")
    combined_df["标的代码"] = pd.to_numeric(combined_df["标的代码"], errors="coerce")
    return combined_df  # type: ignore


def calculate_option_value(df: pd.DataFrame) -> pd.DataFrame:
    # 计算保证金
    df = calculate_option_margin(df)
    # 计算收益率，年化收益率，虚值度
    df = calculate_option_radio(df)

    # 这里再调整顺序
    output_columns = [
        "期权名称", "最新价", "时间价值", "内在价值", "理论价格", "Delta", "虚值度", "收益率", "年化收益率", "剩余天数"
    ]
    rest_columns = [col for col in df.columns if col not in output_columns]
    df = df[output_columns + rest_columns]
    #这里增加一个逻辑 最新价 保留4位小数
    df['最新价'] = df['最新价'].round(4)

    return df


def calculate_option_margin(df: pd.DataFrame, contract_unit: int = 10000,
                            broker_margin_ratio: float = 0.2) -> pd.DataFrame:
    """
    计算期权开仓保证金和维持保证金

    参数:
    df: 期权数据DataFrame，应包含'期权名称'、'最新价'、'标的最新价'列
    contract_unit: 合约单位，默认10000
    broker_margin_ratio: 券商在交易所基础上增加的保证金比例，默认0.2 (20%)

    返回:
    添加了保证金计算结果的DataFrame
    """
    # 从期权名称中提取行权价
    pattern = re.compile(r'^(.+?)([沽购])(\d{1,2})月(\d{3,4})')
    df['行权价'] = df['期权名称'].apply(lambda x: float(pattern.match(x).group(4)) / 1000.0 if pattern.match(x) else None)

    # 判断期权类型(认购/认沽)
    df['期权类型'] = df['期权名称'].apply(lambda x: '认购' if '购' in x else '认沽')

    # 计算虚值额
    df['虚值额'] = 0.0
    # 认购期权虚值额 = max(行权价 - 标的最新价, 0)
    df.loc[df['期权类型'] == '认购', '虚值额'] = df.loc[df['期权类型'] == '认购'].apply(
        lambda row: max(row['行权价'] - row['标的最新价'], 0), axis=1
    )
    # 认沽期权虚值额 = max(标的最新价 - 行权价, 0)
    df.loc[df['期权类型'] == '认沽', '虚值额'] = df.loc[df['期权类型'] == '认沽'].apply(
        lambda row: max(row['标的最新价'] - row['行权价'], 0), axis=1
    )

    # 计算交易所开仓保证金
    df['交易所开仓保证金'] = 0.0
    # 认购期权开仓保证金
    call_margin = df.loc[df['期权类型'] == '认购'].apply(
        lambda row: round((row['最新价'] + max(0.12 * row['标的最新价'] - row['虚值额'], 0.07 * row['标的最新价'])) * contract_unit, 2),
        axis=1
    )
    df.loc[df['期权类型'] == '认购', '交易所开仓保证金'] = call_margin

    # 认沽期权开仓保证金
    put_margin = df.loc[df['期权类型'] == '认沽'].apply(
        lambda row: round(min(
            row['最新价'] + max(0.12 * row['标的最新价'] - row['虚值额'], 0.07 * row['行权价']),
            row['行权价']
        ) * contract_unit, 2),
        axis=1
    )
    df.loc[df['期权类型'] == '认沽', '交易所开仓保证金'] = put_margin

    # 计算交易所维持保证金(假设结算价等于最新价)
    df['交易所维持保证金'] = 0.0
    # 认购期权维持保证金
    call_maintenance = df.loc[df['期权类型'] == '认购'].apply(
        lambda row: round((row['最新价'] + max(0.12 * row['标的最新价'] - row['虚值额'], 0.07 * row['标的最新价'])) * contract_unit, 2),
        axis=1
    )
    df.loc[df['期权类型'] == '认购', '交易所维持保证金'] = call_maintenance

    # 认沽期权维持保证金
    put_maintenance = df.loc[df['期权类型'] == '认沽'].apply(
        lambda row: round(min(
            row['最新价'] + max(0.12 * row['标的最新价'] - row['虚值额'], 0.07 * row['行权价']),
            row['行权价']
        ) * contract_unit, 2),
        axis=1
    )
    df.loc[df['期权类型'] == '认沽', '交易所维持保证金'] = put_maintenance

    # 计算券商保证金 = 交易所保证金 * (1 + 券商增加比例)
    df['券商开仓保证金'] = round(df['交易所开仓保证金'] * (1 + broker_margin_ratio), 2)
    df['券商维持保证金'] = round(df['交易所维持保证金'] * (1 + broker_margin_ratio), 2)

    return df


def calculate_option_radio(df: pd.DataFrame, contract_unit: int = 10000) -> pd.DataFrame:
    # 计算虚值度
    df['虚值度'] = 0.0
    # 认购期权虚值度
    df.loc[df['期权类型'] == '认购', '虚值度'] = df.loc[df['期权类型'] == '认购'].apply(
        lambda row: 0 if row['标的最新价'] == 0 or row['行权价'] == 0 else
        ((row['行权价'] / row['标的最新价']) - 1) * 100,
        axis=1
    )

    # 认沽期权虚值度
    df.loc[df['期权类型'] == '认沽', '虚值度'] = df.loc[df['期权类型'] == '认沽'].apply(
        lambda row: 0 if row['标的最新价'] == 0 or row['行权价'] == 0 else
        ((row['标的最新价'] - row['行权价']) / row['标的最新价']) * 100,
        axis=1
    )
    df['虚值度'] = df['虚值度'].round(2)

    # 计算最大保证金
    df['保证金'] = df[['券商开仓保证金', '券商维持保证金']].max(axis=1)

    # 计算收益率 (%)
    df['收益率'] = round(100 * df['最新价'] * contract_unit / df['保证金'], 2)

    # 计算剩余天数
    today = datetime.now().date()
    df['剩余天数'] = (pd.to_datetime(df['到期日']) - pd.to_datetime(today)).dt.days

    # 计算年化收益率 (%)
    df['年化收益率'] = round(df['收益率'] * 365 / df['剩余天数'], 2)

    # 转换到期日为datetime类型并处理异常
    df['到期日'] = pd.to_datetime(df['到期日'], errors='coerce', format='%Y-%m-%d')
    # 计算剩余天数
    today = pd.Timestamp.now()
    df['剩余天数'] = (df['到期日'] - today).dt.days

    return df


def filter_by_degree(df: pd.DataFrame,
                     virtual_threshold_map: dict = VIR_VALUE_DEGREE_THRESHOLD_DICT,
                     return_threshold_map: dict = ANNUALIZED_RETURN_THRESHOLD_DICT,
                     delta_threshold_map: dict = DELTA_THRESHOLD_DICT) -> pd.DataFrame:
    """
    根据虚值度阈值、年化收益率阈值和期权期限过滤期权数据

    参数:
    df: 期权数据DataFrame，应包含'标的代码'、'虚值度'、'年化收益率'、'期权类型'和'到期日'列
    virtual_threshold_map: 虚值度阈值映射表
    return_threshold_map: 年化收益率阈值映射表

    返回:
    过滤后的DataFrame
    """

    # 创建默认掩码，初始为True（保留所有行）
    mask = pd.Series([True] * len(df), index=df.index)

    # 过滤虚值度<0或者收益率<1的数据
    mask = mask & (df['虚值度'] >= 0) & (df['收益率'] >= 1)
    print(f"过滤虚值度<0或收益率<1的数据后，剩余数据量: {mask.sum()}")

    # 遍历每个标的代码的虚值度阈值
    for underlying_code, threshold in virtual_threshold_map.items():
        underlying_mask = df['标的代码'] == underlying_code
        mask = mask & (~underlying_mask | (df['虚值度'] >= threshold))

    # 遍历每个期权类型的年化收益率阈值
    for option_type, threshold in return_threshold_map.items():
        option_mask = df['期权类型'] == option_type
        mask = mask & (~option_mask | (df['年化收益率'] >= threshold))

    # # 根据期权类型和delta 做过滤
    # for option_type, threshold in delta_threshold_map.items():
    #     delta_mask = df['期权类型'] == option_type
    #     mask = mask & (~delta_mask | (df['Delta'] >= threshold))


    # 过滤无效日期
    invalid_dates = df['到期日'].isna()
    if invalid_dates.sum() > 0:
        print(f"警告: 发现{invalid_dates.sum()}个无效日期，已过滤")
        df = df[~invalid_dates].copy()  # type: ignore
        mask = mask[~invalid_dates]



    # 按期权类型分别处理期限过滤
    call_options_mask = df['期权类型'] == '认购'
    put_options_mask = df['期权类型'] == '认沽'

    # 认购期权：保留最近两个月
    if call_options_mask.sum() > 0:
        min_days = df.loc[call_options_mask, '剩余天数'].min()
        call_keep_mask = call_options_mask & (df['剩余天数'] <= min_days + 60)
        mask = mask & (call_keep_mask | ~call_options_mask)

    # 认沽期权：保留最远月
    if put_options_mask.sum() > 0:
        max_days = df.loc[put_options_mask, '剩余天数'].max()
        put_keep_mask = put_options_mask & (df['剩余天数'] == max_days)
        mask = mask & (put_keep_mask | ~put_options_mask)

    # 应用掩码过滤数据
    filtered_df = df[mask].copy()
    print(f"过滤前数据量: {len(df)}, 过滤后数据量: {len(filtered_df)}")

    return filtered_df  # type: ignore


def filter_by_option_info(df: pd.DataFrame, query: Union[str, List[str]]) -> pd.DataFrame:
    """
    根据输入字符串或字符串数组自动选择过滤列：
    - 如果字符串是数字，则匹配期权代码列
    - 如果字符串是文本，则匹配期权名称列
    - 如果输入是数组，则对每个元素执行上述逻辑，结果合并

    参数:
    df: 期权数据DataFrame，应包含'期权代码'和'期权名称'列
    query: 输入的查询字符串或字符串数组

    返回:
    过滤后的DataFrame
    """
    # 如果输入是字符串，转换为列表
    if isinstance(query, str):
        query_list = [query]
    else:
        query_list = query

    # 存储所有过滤结果
    all_filtered_dfs = []

    for single_query in query_list:
        try:
            # 尝试将字符串转换为数字
            option_code = int(single_query)
            # 如果转换成功，按期权代码列过滤
            filtered_df = df[df['期权代码'] == option_code].copy()
            print(f"按期权代码过滤: {option_code}, 匹配结果: {len(filtered_df)} 条")
            all_filtered_dfs.append(filtered_df)
        except ValueError:
            # 如果转换失败，按期权名称列过滤（包含匹配）
            filtered_df = df[df['期权名称'].str.contains(single_query, na=False)].copy()
            print(f"按期权名称过滤: {single_query}, 匹配结果: {len(filtered_df)} 条")
            all_filtered_dfs.append(filtered_df)

    # 合并所有结果
    if all_filtered_dfs:
        result_df = pd.concat(all_filtered_dfs, ignore_index=True)
        # 去重
        result_df = result_df.drop_duplicates()
        print(f"合并后总结果: {len(result_df)} 条")
        return result_df
    else:
        # 如果没有匹配结果，返回空的DataFrame
        return pd.DataFrame(columns=df.columns)


def do_analysis() -> pd.DataFrame:
    df = fetch_option_risk_value()
    df = calculate_option_value(df)
    df = filter_by_degree(df)

    return df  # type: ignore


def do_calc_filter(query: Union[str, List[str]]):
    df = fetch_option_risk_value()
    df = calculate_option_value(df)
    df = filter_by_option_info(df=df, query=query)
    return df


def filter_far_put_lowest_strike_top2_by_return(df: pd.DataFrame) -> pd.DataFrame:
    """
    1. 只保留最远月的认沽期权
    2. 每一个标的的期权，只获取行权价格最低的2个期权
    3. 结果按照收益率逆序排列
    """
    # 只保留认购期权
    call_df = df[df['期权类型'] == '认沽'].copy()
    if call_df.empty:
        return call_df

    # 只保留最远月
    max_expiry = call_df['到期日'].max()
    far_month_df = call_df[call_df['到期日'] == max_expiry].copy()
    if far_month_df.empty:
        return far_month_df

    # 从期权名称中提取行权价（如未提取过）
    if '行权价' not in far_month_df.columns:
        pattern = re.compile(r'^(.+?)([沽购])(?\d{1,2})月(\d{3,4})')
        far_month_df['行权价'] = far_month_df['期权名称'].apply(
            lambda x: float(pattern.match(x).group(4)) / 1000.0 if pattern.match(x) else None
        )

    # 每个标的只保留行权价最低的2个期权
    filtered_df = (
        far_month_df
            .sort_values('行权价')
            .groupby('标的代码', group_keys=False)
            .head(2)
    )

    # 按收益率逆序排列
    filtered_df = filtered_df.sort_values(by='收益率', ascending=False)
    filtered_df = filtered_df.reset_index(drop=True)
    return filtered_df


def filter_near_call_highest_strike_top2_by_return(df: pd.DataFrame) -> pd.DataFrame:
    """
    1. 只保留本月和下个月的认购期权
    2. 每一个标的的期权，只获取行权价格最高的2个期权
    3. 结果按照虚值度逆序排列
    """
    # 只保留认购期权
    call_df = df[df['期权类型'] == '认购'].copy()
    if call_df.empty:
        return call_df

    # 获取本月和下个月的到期日
    unique_expiries = sorted(call_df['到期日'].unique())
    if len(unique_expiries) < 2:
        # 如果只有一个月的数据，就只保留那个月
        near_months_df = call_df[call_df['到期日'] == unique_expiries[0]].copy()
    else:
        # 保留最近的两个月
        near_months = unique_expiries[:2]
        near_months_df = call_df[call_df['到期日'].isin(near_months)].copy()

    if near_months_df.empty:
        return near_months_df

    # 从期权名称中提取行权价（如未提取过）
    if '行权价' not in near_months_df.columns:
        pattern = re.compile(r'^(.+?)([沽购])(?\d{1,2})月(\d{3,4})')
        near_months_df['行权价'] = near_months_df['期权名称'].apply(
            lambda x: float(pattern.match(x).group(4)) / 1000.0 if pattern.match(x) else None
        )

    # 每个标的只保留行权价最高的2个期权
    filtered_df = (
        near_months_df
            .sort_values('行权价', ascending=False)  # 按行权价降序排列
            .groupby('标的代码', group_keys=False)
            .head(2)
    )

    # 按收益率逆序排列
    filtered_df = filtered_df.sort_values(by='虚值度', ascending=False)
    filtered_df = filtered_df.reset_index(drop=True)
    return filtered_df

def filter_month_type_lowest_strike_top2_by_return(df: pd.DataFrame, month: str, option_type: str) -> pd.DataFrame:
    """
    只保留指定月份的认沽或认购期权，每个标的只获取行权价最高|最低的2个期权，结果按收益率逆序排列
    # 方法功能调整：认购期权取行权价最高2个，认沽期权取行权价最低2个
    # 只保留指定类型的期权
    """
    filtered_df = df[df['期权类型'] == option_type].copy()
    if filtered_df.empty:
        return filtered_df

    # 只保留指定月份的期权
    def month_match(x):
        if isinstance(x, str):
            return (month in x) or (month.replace("-", "/") in x) or (month.replace("/", "-") in x) or (month.replace("月", "") in x)
        elif hasattr(x, 'strftime'):
            return month in x.strftime("%Y-%m") or month in x.strftime("%Y/%m") or month.replace("月", "") in x.strftime("%m")
        return False

    filtered_df = filtered_df[filtered_df['到期日'].apply(month_match)].copy()
    if filtered_df.empty:
        return filtered_df

    # 从期权名称中提取行权价（如未提取过）
    if '行权价' not in filtered_df.columns:
        pattern = re.compile(r'^(.+?)([沽购])(?\d{1,2})月(\d{3,4})')
        filtered_df['行权价'] = filtered_df['期权名称'].apply(
            lambda x: float(pattern.match(x).group(4)) / 1000.0 if pattern.match(x) else None
        )

    # 根据期权类型决定排序方式
    if option_type == "认购":
        # 认购期权取行权价最高的2个
        filtered_df = (
            filtered_df
                .sort_values('行权价', ascending=False)
                .groupby('标的代码', group_keys=False)
                .head(2)
        )
    else:
        # 认沽期权取行权价最低的2个
        filtered_df = (
            filtered_df
                .sort_values('行权价', ascending=True)
                .groupby('标的代码', group_keys=False)
                .head(2)
        )

    # 按收益率逆序排列
    filtered_df = filtered_df.sort_values(by='收益率', ascending=False)
    filtered_df = filtered_df.reset_index(drop=True)
    return filtered_df

if __name__ == "__main__":
    # df = do_analysis()
    # price_cols = [col for col in df.columns if "价" in col]
    # for col in price_cols:
    #     df[col] = df[col].apply(lambda x: f"{x:.4f}" if pd.notnull(x) else "")
    # df.to_csv("../result/option_risk_value_filter_degree.csv", index=False, encoding='utf-8-sig')

    # # 数组查询示例
    # df_array = do_calc_filter(["科创50购9月", "科创板50购9月"])
    # # df_array = do_calc_filter(["科创50沽8月","科创50沽9月", "科创50沽12月"])
    # # df_array = do_calc_filter(["300ETF沽8月", "300ETF沽9月", "300ETF沽12月"])
    # # df_array = do_calc_filter(["沽9月"])
    #
    # # 过滤虚值度<0或者收益率<1的数据
    # df_array = df_array[(df_array['虚值度'] >= 0) & (df_array['收益率'] >= 1)]
    # print(f"过滤虚值度<0或收益率<1的数据后，剩余数据量: {len(df_array)}")
    #
    # # 对所有包含"价"字的列，保留4位小数，补零，转为字符串
    # price_cols = [col for col in df_array.columns if "价" in col]
    # for col in price_cols:
    #     df_array[col] = df_array[col].apply(lambda x: f"{x:.4f}" if pd.notnull(x) else "")
    # df_array.to_csv("../result/option_risk_value_filter_query_array.csv", index=False, encoding='utf-8-sig')

    # 最远月认沽期权
    # df = fetch_option_risk_value()
    # df = calculate_option_value(df)
    # df = filter_far_put_lowest_strike_top2_by_return(df)
    # price_cols = [col for col in df.columns if "价" in col]
    # for col in price_cols:
    #     df[col] = df[col].apply(lambda x: f"{x:.4f}" if pd.notnull(x) else "")
    # df.to_csv("../result/filter_far_put_lowest_strike_top2.csv", index=False, encoding='utf-8-sig')

    # 指定月份的认沽期权
    df_fetch = fetch_option_risk_value()
    df = calculate_option_value(df_fetch)
    month = "6月"
    option_type = "认购"
    df = filter_month_type_lowest_strike_top2_by_return(df, month, option_type)
    price_cols = [col for col in df.columns if "价" in col]
    for col in price_cols:
        df[col] = df[col].apply(lambda x: f"{x:.4f}" if pd.notnull(x) else "")
    csv_file = f"../result/{month}_{option_type}_lowest_strike_top2.csv"
    df.to_csv(csv_file, index=False, encoding='utf-8-sig')
    print(csv_file)

    df = calculate_option_value(df_fetch)
    month = "6月"
    option_type = "认沽"
    df = filter_month_type_lowest_strike_top2_by_return(df, month, option_type)
    price_cols = [col for col in df.columns if "价" in col]
    for col in price_cols:
        df[col] = df[col].apply(lambda x: f"{x:.4f}" if pd.notnull(x) else "")
    csv_file = f"../result/{month}_{option_type}_lowest_strike_top2.csv"
    df.to_csv(csv_file, index=False, encoding='utf-8-sig')
    print(csv_file)

    # 本月和下月认购期权（行权价最高的2个）
    # df_call = fetch_option_risk_value()
    # df_call = calculate_option_value(df_call)
    # df_call = filter_near_call_highest_strike_top2_by_return(df_call)
    # price_cols = [col for col in df_call.columns if "价" in col]
    # for col in price_cols:
    #     df_call[col] = df_call[col].apply(lambda x: f"{x:.4f}" if pd.notnull(x) else "")
    # df_call.to_csv("../result/filter_near_call_highest_strike_top2.csv", index=False, encoding='utf-8-sig')
